Lower mean.dim beyond 4D pooling as an XNNPACK static mean by JakeStevens · Pull Request #24741 · pytorch/executorch

JakeStevens · 2026-11-06T18:03:40Z

Ulterior Motive

Run RMSNorm in LeRobot policies on XNNPACK (#24667), where mean.dim fails to partition with "only supports averaging 4D tensors, got tensor of rank 3".

Rationale

What: Add an XNNStaticMean node backed by xnn_define_static_reduce(xnn_reduce_mean), and lower every mean.dim that is not Global Average Pooling to it: ranks 1 to 6, any dims, and either keepdim.

Why: XNNPACK only lowered a 4D mean over dims [2, 3] with keepdim=True. RMSNorm averages the last dim of a 3-D tensor, so its mean stayed on the portable kernel and split the delegate around every norm.

Details

Review order: the serialization schema and XNNCompiler.cpp, then op_mean_dim.py and is_global_avg_pool_mean, then ChannelsLastTaggedReshapePass, then MeanDimConfig, then the tests.

The 4D, dims [2, 3], keepdim=True case still lowers to Global Average Pooling on NHWC. For those programs the .pte is byte-identical to main (fp32, fp16, qs8, conv-bn-hardtanh-mean, and a CNN ending in AdaptiveAvgPool2d(1)), and the tests check which XNNPACK node each mean serializes to and how many transposes surround it.
Every other mean.dim reduces over NCHW dims, so ChannelsLastTaggedReshapePass gives it an NCHW input, and the visitor refuses an NHWC one as an internal layout error. After a convolution that adds layout copies inside the one delegate. For conv then mean(1, keepdim=True) then conv, the program has 4 transposes, as on main. For y - y.mean(1, keepdim=True) after a convolution (the centering in ConvNeXt's LayerNorm2d), the subtraction stays NHWC, so its full-size result is copied back to NCHW at the output: 4 transposes in one delegate, where main has 2 in two delegates around a portable mean. At 96 channels and 56x56 that is 0.75x at 1 thread, 0.85x at 4 and 1.12x at 16, and conv, mean, conv is 0.54x, 0.70x and 0.90x. Remapping the reduction axes to read NHWC directly would remove these copies, also for heads like x.mean([2, 3]); that is left for a follow-up.
is_global_avg_pool_mean takes the rank from the mean's own output, which with keepdim=True equals the input's, because inside ChannelsLastTaggedReshapePass the mean's input can be a newly inserted node without "val".
The node uses a shared _XNNStaticReduce table, so other XNNPACK reductions can reuse it and defineGenericReduceNode. It is appended to XNodeUnion, so existing node IDs do not change. Practically every re-exported XNNPACK LLM program contains it, because RMSNorm (examples/models/llama/norm.py) and QK-norm lower to it, and those programs fail to load on runtimes built before this change, including released prebuilt runtimes, with Unhandled node type in the log. The release note should say so.
ChannelsLastTaggedReshapePass now reads the layout of a quantized graph input from the input's own tag. It used to test meta["val"][0], a 3-D slice that is never channels_last, so a channels_last input behind a quantize and dequantize read as NCHW. A mean over NCHW dims of such an input then reduced the wrong axes, and on main a convolution of it got a transpose that fails at runtime. Both now lower correctly.
An empty dim list averages over every dim, as in PyTorch.
Latency against main, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: RMSNorm 1x128x2048 0.48x, 3-D last-dim mean 0.45x, RMSNorm then Linear 0.77x, 4 pre-norm residual MLP blocks 0.99x for decode (1 token) and 0.80x for prefill (64 tokens), conv then mean over channels 0.88x.
XNNPACK splits a reduction with more than one output into threadpool tasks, so a small mean with several outputs can cost more in XNNPACK than on the portable kernel when the pool has many threads. For RMSNorm on 10 shapes from 1x1x128 to 1x32x2048, every shape is 0.43x to 0.97x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.55x (1.5 to 2.7 us more), and the rest 0.44x to 0.93x. At 16 threads those small shapes are 3.7x to 5.7x (21 to 23 us more), and the rest 0.70x to 0.98x. Inside a model with neighbouring delegated ops the cost can remain: decode QK-norm between two linears, 8 heads of 64, goes from 11.8 us to 39.1 us at 16 threads, and 16 heads of 128 from 222 us to 241 us. A mean alone in its delegate pays the same wake-up (a bare 1x128x128 mean is 3.26x at 16 threads and 1.14x at 4), as does any op lowered alone, such as add (2.95x at 1x128x128 and 16 threads on main). These means are lowered like every other op the XNNPACK partitioner takes, rather than gated by size, which would be tuned to one thread count and treat static and dynamic exports differently. End to end, Qwen3-0.6B with 8da4w quantization (113 RMSNorm and QK-norm means, 368 delegates on main and 283 here) decodes, after over before in executor_runner with the LLM custom ops: on the same Apple M-series CPU 0.98x at 1 thread and 1.11x to 1.12x at the default 16 threads (97 to 87 tokens/s); on an iPhone 17 1.00x at 1 thread and 1.05x at the runner's default (101 to 96 tokens/s); on a Pixel 10 0.98x and 1.00x. The cost is the multi-head QK-norm means. Running small multi-output reductions without waking the pool belongs in XNNPACK and is left for a follow-up.
The LoRA CI scripts check the start of the 8da4w-quantized Qwen3 base model's greedy output with a regex. With RMSNorm and QK-norm now in XNNPACK, its first token can tip from :\n to : (assistant: me instead of assistant:\n<think>): unquantized, main and this branch generate the same tokens with logits within 1e-4, but 8da4w re-rounds activations and the two candidates are within about 0.5 logits. The regex accepts exactly those two continuations, and the scripts print it when it does not match.
Tests cover 1-D to 6-D inputs, all dims, keepdim=False, fp16 and qs8 over the innermost and a middle dim, dynamic kept and reduced dims over 5 random shapes, RMSNorm, fp32 and qs8 convolutions followed by a mean, a mean over channels whose result meets the convolution output or feeds another convolution (as in ConvNeXt's LayerNorm2d), qs8 channels_last inputs to a mean and to a convolution, the serialized XNNPACK nodes (moved xnnpack_node_types from the layout pass tests to the XNNPACK Tester), and the layout copies the pass inserts.

Test plan:

python -m pytest backends/xnnpack/test/ops/test_mean_dim.py backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py
python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer

Authored with AI using Claude Code.

pytorch-bot · 2026-11-06T18:03:46Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24741

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit e578110 with merge base bdc7c35 ():

NEW FAILURE - The following job has failed:

Test QNN Windows MSVC build / build-qnn-windows-x64 / windows-job (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

JakeStevens · 2026-11-06T20:24:11Z

@claude review effort=max

claude · 2026-11-06T20:24:52Z

Claude finished @JakeStevens's task in 36m 9s —— View job

Needs Discussion: the 8192-element cut-off in MeanDimConfig doesn't follow from the PR's own latency data, and it still splits the delegate around small static RMSNorm and QK-norm means; the other findings are test gaps and minor items.

Detailed review

PR Review: #24741

Summary

This PR adds an XNNStaticMean node backed by xnn_define_static_reduce. Every mean.dim other than a 4D [2, 3]/keepdim=True Global Average Pooling now lowers to it on NCHW, except small static means with more than one output. The cut-off for those small means needs discussion. The rest are gaps in the tests and minor cleanups.

Code Quality

backends/xnnpack/utils/utils.py:96 — is_global_avg_pool_mean reads the input's meta["val"] partway through ChannelsLastTaggedReshapePass (channels_last_tagged_reshape_pass.py:320, :329).
- Crash when the conv is visited first. Take a static per-tensor dq shared by a conv and a mean. The conv's input_to_nhwc calls insert_copy_q_dq. Its before.replace_all_uses_with(dq) (:405) moves every user of the old dq onto the new dq, including the mean. The new dq copies its meta from the _to_copy created at :591-598, which has no "val". When the pass then reaches the mean, it raises KeyError. On main, mean.dim was a plain membership check in memory_sensitive_ops_nhwc, so this graph worked for GAP means. If the graph is reachable, this is a regression.
- Silent wrong axes when the mean is visited first. A non-GAP mean is processed, then moved onto the NHWC dq. node_visitor.py:427 ORs in the NHWC tag, so the NCHW reduction_axes index NHWC data. The re-converge at :745-766 only covers broadcast binary ops; view_copy already has the same exposure.
- Confidence is low on reachability. PT2E's DuplicateDQPass normally gives each user its own dq.
- Cheap partial fix. Check keepdim first, then take the rank from the mean's own node.meta["val"]. With keepdim=True it equals the input rank. That removes the KeyError half. Fix this →

Testing

No test asserts which XNNPACK node is serialized.
- test_mean_dim.py only checks that mean.dim is consumed and that the outputs match. GAP and XNNStaticMean produce the same numbers.
- So if the [2, 3]/keepdim=True case were routed through XNNStaticMean, with an extra transpose after a conv, every test would still pass. The description's claim that those .ptes are byte-identical to main would then be false without anything noticing.
- _xnnpack_node_types (test_channels_last_tagged_reshape.py:493) already decompiles the delegate. Move it into a shared helper and assert the XNNGlobalAvgPooling2d, XNNStaticMean and XNNStaticTranspose counts. Fix this →
The dynamic-shape tests run a single random shape.
- test_mean_dim.py:188-196 go through _test_mean_dim, which calls run_method_and_compare_outputs() with the default num_runs=1 (backends/test/harness/tester.py:320).
- Each run therefore draws one seq in [1, 32]. The bounds, including a reduced dim of size 1, are only hit by chance.
- test_add.py, test_softmax.py and test_conv1d.py pass num_runs= for this. Fix this →
Some cases are missing.
- A qs8 conv followed by a non-GAP mean is the only path where the new NCHW requirement meets the quantized copy wrapper (insert_copy_q_dq). test_fp32_conv_mean_dim is fp32-only, and the qs8 tests (:241-251) have no conv and all use keepdim=True.
- The description says ranks 1 to 6 are supported, but the highest rank tested is 5-D (:156).
- Fix this →
The base-model check in .ci/scripts/test_lora.sh:52-53 and .ci/scripts/test_lora_multimethod.sh:52-53 is wider than the observed flip.
- The regex went from assistant:\n(?:<think>\n)+ to assistant:[^\n]*\n(?:<think>\n)*. It now accepts any text after assistant:, and also no <think> at all. The one observed flip is assistant: me.
- Pin the alternate first line, for example assistant:(?: me)?\n. Keep + unless the flipped output really drops <think>.
- On failure, the scripts still print only EXPECTED_QUANT_PREFIX (test_lora.sh:243,247) and EXPECTED_BASE_PREFIX (test_lora_multimethod.sh:156). Someone debugging a failure doesn't see what the check actually accepts. Fix this →

Backward Compatibility

Practically every re-exported Llama-family .pte will now contain XNNStaticMean.
- examples/models/llama/norm.py:35 is (x * x).mean(-1, keepdim=True). Decode RMSNorm has a single output, and dynamic-shape exports skip the cut-off.
- Runtimes built before this PR, such as the released AAR and xcframework, reject the program in defineNotImplementedNode (XNNCompiler.cpp:1831).
- The description does note that an updated runtime is needed. But "such as RMSNorm" undersells it: effectively every new XNNPACK LLM export is affected.
- There is no partitioner option to keep GAP-only lowering. Removing MeanDimConfig would also drop GAP.
- Strictly this is forward rather than backward compatibility, but it is the change users will notice most, and the release note should say so explicitly.

Performance

The min_multi_row_input_numel = 8192 cut-off (generic_node_configs.py:707, :735-745) doesn't follow from the PR's own numbers.
- Measured on a desktop thread count. It was derived at 16 threads on an Apple M-series CPU. At 1 thread the description measures the 1x8x64 case as a tie (~5 µs), and at 4 threads only 1x128x128 is slower. At mobile thread counts it rejects means that cost nothing extra in XNNPACK.
- Slow means above it still lower. The description measured lone means at 1x128x128 (16384 elements, 3.26×) and 1x64x512 (1.90×) as slower, and both are still lowered.
- In-model means below it are rejected. The description itself says that inside a model "the neighbouring delegated ops already pay the wake-up". Yet two examples stay on portable and split the delegate between linears:
  - a static-shape RMSNorm on 1x16x256 (4096 elements, 16 outputs);
  - Qwen3-0.6B's static-decode QK-norm (16×128 and 8×128 per token).
- Check the motivating shapes. It is worth confirming that the LeRobot shapes behind XNNPACK: LeRobot policies that do not run yet, and what blocks them #24667 clear the cut-off.
- Static and dynamic exports diverge. is_shape_dynamic bypasses the check (:732), so the same module partitions differently under static and dynamic export.
- No precedent, and it can't be tuned. No other XNNPACK config gates on size for performance. It is a class attribute rather than an XNNPartitionerConfig option like force_non_static_weights_for_f32_linear (xnnpack_config.py:64), so it can't be set per deployment.
I suggest dropping it and treating a lone small mean as the "existing single-op trade-off" the description already accepts for add. If it stays, justify 8192 with in-model numbers at the shapes it rejects, and make it a partitioner option.
Every non-GAP 4D mean on an NHWC-tagged tensor gets a full NHWC→NCHW copy inside the delegate (channels_last_tagged_reshape_pass.py:328-329).
- Examples are the torchvision ShuffleNetV2 and MNASNet heads (x.mean([2, 3])) and CBAM-style x.mean(1, keepdim=True).
- test_fp32_channels_last_tagged_reshape_pass_conv_mean_channels (test_channels_last_tagged_reshape.py:930-938) locks this in by expecting two copies.
- This isn't a regression, because main paid the same transpose at the delegate boundary.
- Remapping the reduction axes for an NHWC input (NCHW dim 1→3, 2→1, 3→2) would avoid the copy. That works for keepdim=True, and for keepdim=False whenever the kept dims stay in order, which covers both examples. This is fine as a follow-up.

Specific Comments

backends/xnnpack/partition/config/generic_node_configs.py:707 — if the cut-off stays, "multi_row" is not what the condition tests (output_numel > 1). The why message at :743 ("needs at least 8192 input elements") reads like an XNNPACK capability limit rather than a performance heuristic, which will mislead anyone chasing a partition split.
docs/source/backends/xnnpack/op-support.csv:32 — "a static mean" reads like XNNPACK's static-mean node. "A mean with a static shape", as the MeanDimConfig docstring puts it, is clearer. The hard-coded 8192 will also drift from min_multi_row_input_numel.
.ci/scripts/test_lora.sh:49-50 (same in test_lora_multimethod.sh) — "can flip between builds" hides the cause. Per the description, the flip comes from RMSNorm and QK-norm moving into XNNPACK under 8da4w.
backends/xnnpack/runtime/XNNCompiler.cpp:2059 — the in-tree CPU provider calls xnn_define_static_reduce_v2 (with int64_t axes) against the same XNNPACK pin (backends/cpu/runtime/providers/xnnpack/XNNPACKProvider.cpp:633). If v1 is deprecated there, -Wno-deprecated-declarations (backends/xnnpack/CMakeLists.txt:40) would hide the warning. Confidence is low, because I couldn't read the pinned header.

Recommendation

Needs Discussion

The cut-off is the open design question: either drop it, or justify it with in-model data and make it a partitioner option. The node-type assertions and the tighter LoRA regex should land with this PR. Everything else is minor or can be a follow-up.

### Ulterior Motive Improving overall BF16 support in light of #24667 ### Summary `ToDimOrderCopyConfig` already lowers fp32 <-> fp16 casts to XNNPACK convert nodes. This allows fp32 <-> bf16 as well, when bf16 is enabled in the partitioner. fp16 <-> bf16 stays unsupported, because XNNPACK has no direct convert between them. bf16 models that upcast a step to fp32 (HF softmax with `dtype=torch.float32`, Llama-style RMSNorm) now stay in one delegate, instead of running the casts in portable kernels between partitions. ExecuTorch's current XNNPACK pin already has the bf16 <-> fp32 convert kernels, so no runtime or pin changes are needed. The config matches `_to_dim_order_copy`, which export emits when dim order is kept, the `to_edge_transform_and_lower` default. Programs exported with `skip_dim_order=True`, as `LLMEdgeManager` and `export_llama` do by default, keep `aten._to_copy`, which no XNNPACK partitioner config matches yet, so their casts (fp16 ones included) still run in portable kernels. Partitioning `_to_copy` dtype casts under the same rules would be a follow-up. ### Benchmarks Measured with `executor_runner` (Release; XNNPACK + portable ops), median of 50 iterations. "Before" exports the same model with the cast lowering disabled. Shapes: - softmax: `softmax(x, -1, dtype=float32).to(bf16)` on `[1,16,S,S]`. - RMSNorm: Llama-style, `[1,T,2048]`. - elementwise: `(x.float()*2+1).to(bf16)` on `[1,128,4096]`. Single thread (Android pinned to the fastest core with `taskset`), ms, before -> after: | model | Apple M4 | iPhone 17 (A19) | Pixel 10 (Tensor G5, X4) | Pixel 6 (Tensor G1, X1; no BF16) | |---|---|---|---|---| | softmax S=128 | 1.172 -> 0.141 (8.3x) | 1.415 -> 0.403 (3.5x) | 1.785 -> 0.701 (2.5x) | 1.760 -> 0.684 (2.6x) | | softmax S=512 | 18.64 -> 2.343 (8.0x) | 20.18 -> 3.246 (6.2x) | 31.16 -> 5.031 (6.2x) | 35.60 -> 6.484 (5.5x) | | RMSNorm T=1 | 0.012 -> 0.002 (5.6x) | 0.021 -> 0.004 (4.7x) | 0.098 -> 0.012 (8.3x) | 0.013 -> 0.004 (3.1x) | | RMSNorm T=128 | 1.266 -> 0.243 (5.2x) | 1.620 -> 0.652 (2.5x) | 2.502 -> 0.820 (3.1x) | 1.731 -> 0.712 (2.4x) | | elementwise | 2.161 -> 0.118 (18.3x) | 2.244 -> 0.418 (5.4x) | 3.754 -> 0.467 (8.0x) | 4.053 -> 0.576 (7.0x) | Default `executor_runner` threads (M4: 16, iPhone: 6, Pixel 10: 6, Pixel 6: 4; unpinned), ms, before -> after: | model | Apple M4 | iPhone 17 (A19) | Pixel 10 (Tensor G5) | Pixel 6 (Tensor G1) | |---|---|---|---|---| | softmax S=128 | 1.199 -> 0.152 (7.9x) | 1.483 -> 0.294 (5.0x) | 3.858 -> 1.020 (3.8x) | 1.980 -> 0.317 (6.2x) | | softmax S=512 | 17.11 -> 0.635 (27x) | 18.98 -> 1.956 (9.7x) | 38.88 -> 5.518 (7.0x) | 25.97 -> 5.074 (5.1x) | | RMSNorm T=1 | 0.010 -> 0.002 (4.8x) | 0.020 -> 0.004 (4.4x) | 0.101 -> 0.101 (1.0x)* | 0.133 -> 0.025 (5.3x) | | RMSNorm T=128 | 1.323 -> 0.347 (3.8x) | 1.616 -> 0.720 (2.2x) | 8.936 -> 1.765 (5.1x) | 2.544 -> 0.642 (4.0x) | | elementwise | 2.154 -> 0.154 (14x) | 2.376 -> 0.355 (6.7x) | 3.241 -> 0.839 (3.9x) | 2.806 -> 0.462 (6.1x) | \* Unpinned, this ~10 us graph runs on a core that doesn't ramp its clock up, so the run hits a ~0.1 ms floor even with 1 thread. Pinned, it's 0.012 ms (table above). RMSNorm still has `mean` outside the delegate (`mean.dim` over the last dim isn't lowered yet; #24741 lowers it). The casts no longer split it further (3 -> 2 delegates). ### Test plan ``` python -m pytest backends/xnnpack/test/ops/test_to_copy_dtype.py backends/xnnpack/test/ops/test_to_copy.py ``` New tests cover: - fp32 I/O around a bf16 linear; - a bf16 upcast; - a bf16 softmax computed in fp32 and cast back, as one delegate; - fp32 -> bf16 rounding matching PyTorch exactly (zero tolerance), including values halfway between two bf16 numbers; - no lowering without `enable_bf16`; - fp16 <-> bf16 not partitioned. The tests that run bf16 in XNNPACK are skipped in fbcode, like the other XNNPACK bf16 runtime tests. The delegation tests fail without the config change.

## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` as a convert to fp32, the fp32 `XNNStaticMean` from pytorch#24741, and a convert back to bf16, all inside the delegate. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - The converts are written into the program rather than added by the runtime at load time, as bf16 matmul results are. That keeps the change in Python and matches eager bf16 exactly, but a program exported now keeps the converts if XNNPACK's `xnn_define_static_reduce` later accepts bf16, which is the only thing that rejects a bf16 mean today. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` goes through the same converts as any other bf16 mean. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` as a convert to fp32, the fp32 `XNNStaticMean` from pytorch#24741, and a convert back to bf16, all inside the delegate. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - The converts are written into the program rather than added by the runtime at load time, as bf16 matmul results are. That keeps the change in Python and matches eager bf16 exactly, but a program exported now keeps the converts if XNNPACK's `xnn_define_static_reduce` later accepts bf16, which is the only thing that rejects a bf16 mean today. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` goes through the same converts as any other bf16 mean. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

JakeStevens · 2026-11-09T17:24:01Z

@claude review effort=max

claude · 2026-11-09T17:24:44Z

Claude finished @JakeStevens's task in 38m 54s —— View job

Request Changes: a quantized channels_last input that feeds a mean other than Global Average Pooling now fails to lower. Two things need discussion: default LLM exports now emit a node that older runtimes reject, with no opt-out, and there is no end-to-end decode number for QK-norm. The rest are test gaps and nits.

Detailed review

PR Review: #24741

Summary

This PR lowers every mean.dim other than Global Average Pooling (GAP) to a new XNNStaticMean node on the NCHW input. GAP stays the lowering for 4-D means over [2, 3] with keepdim. Problems:

One case that exports on main now fails to lower.
The new node has no opt-out for runtimes that predate it.
The QK-norm decode cost has no end-to-end measurement.
The tests leave the NHWC interaction and some node counts unchecked.

I could not run Python in this job, and the XNNPACK submodule was not checked out. Every finding comes from reading the code.

Backward Compatibility

A quantized model whose channels_last input feeds a non-GAP mean now fails to lower. I traced this through the code but did not run it. Take x (4-D, channels_last) → q → dq → mean.dim(dq, [1], keepdim=True):
1. call() tags the placeholder NHWC (channels_last_tagged_reshape_pass.py:676-677). The else branch (:726-728) then tags q and dq NHWC.
2. requires_nchw_inputs sends the mean to input_to_nchw(dq), which returns at :636-637 without inserting a copy. The reason is that is_nchw_node(dq) looks through to the placeholder and tests actual_node.meta["val"][0] (:157-160), which is a 3-D slice. is_contiguous(memory_format=torch.channels_last) is False for every 3-D tensor, so the look-through always reports NCHW.
3. dq keeps its NHWC tag, and the new guard (op_mean_dim.py:73-78) raises mean.dim reduces over NCHW dims, but its input is NHWC inside to_edge_transform_and_lower.
On main, MeanDimConfig rejected this mean, so it ran on the portable kernel and the export succeeded. The guard is right to refuse, because letting it through would reduce over the wrong axes. The fix belongs in the look-through in is_nhwc_node and is_nchw_node (:141-143, :158-160). It should read the placeholder's XNN_NHWC_NODE tag, or meta["val"] without the [0].

That fix also changes how conv and the other ops treat a dequantized channels_last input, and no test currently gives the pass a quantized channels_last input. Add qs8 tests for this mean (e.g. MeanDim((1,)) on torch.randn(1, 5, 4, 4).to(memory_format=torch.channels_last)) and for a conv. fp32 is not affected, because input_to_nchw checks placeholders with _is_nhwc(input_node.meta["val"]) (:633-635).
Fix this →
Default LLM exports now emit XNNStaticMean, older runtimes reject it, and there is no way to keep the old lowering.
- Who is affected. extended_ops defaults to True (extension/llm/export/config/llm_config.py:542), so export_llama_lib.py:1676-1681 adds the full XnnpackPartitioner. RMSNorm (examples/models/llama/norm.py:35,53) and QK-norm are last-dim mean(-1, keepdim=True). Every XNNPACK LLM .pte exported with the default settings now contains the new node. The LoRA CI regex change in this PR is a symptom of that.
- What fails. A runtime without this PR fails to load the method in defineNotImplementedNode (XNNCompiler.cpp:1838-1842). The log reads Unhandled node type: with no name, because the old flatc-generated EnumNameXNodeUnion returns "" for an unknown value. That comes from flatc's codegen pattern; I did not inspect the generated header.
- No way out. The only way to avoid the node is to leave MeanDimConfig out of configs=, which also loses GAP. A second mean config is not possible either, because ConfigerationBasedPartitioner raises on a duplicate target name (exir/backend/canonical_partitioners/config_partitioner.py:156-160).
- Suggested fix. ExecuTorch doesn't promise forward compatibility, but this PR changes the lowering of an op that every LLM uses. Add a MeanDimConfig kwarg that keeps only GAP, passed through XnnpackPartitioner(**kwargs) the way force_non_static_weights_for_f32_linear and enable_bf16 are (xnnpack_config.py:60-71). Also name the Unhandled node type symptom in the release notes.
  Fix this →

Testing

No test covers a non-GAP mean whose output meets NHWC data or feeds an NHWC op. test_fp32_conv_mean_dim and test_qs8_conv_mean_dim only cover conv → mean → output. This PR now delegates the LayerNorm2d pattern (ConvNeXt), which has two untested paths:
- In y - y.mean(1, keepdim=True) after a conv, the untagged mean output meets the NHWC conv output in converge_layouts.
- In conv → mean → conv, input_to_nhwc has to copy the mean output back to NHWC.
Add fp32 and qs8 tests for both.
Fix this →
Some node-count checks are missing.
- test_fp32_rms_norm (test_mean_dim.py:202) only checks for one delegate and no edge mean.dim. It still passes if pow or the last mul falls back to portable. Add check_xnnpack_node_count({"XNNStaticMean": 1, "XNNGlobalAvgPooling2d": 0}) and a check_not for the edge pow, rsqrt, add and mul ops.
- _test_qs8_mean_dim and test_qs8_conv_mean_dim don't count XNNConvert, so they can't tell a qs8 mean from an fp32 mean between converts. For _test_qs8_mean_dim the count should be 2: the input quantize and the output dequantize.
  Fix this →

Performance

There is no end-to-end decode number for a model with QK-norm. The description measures decode QK-norm between two linears at 16 threads: 8 heads of 64 goes from 11.8 µs to 39.1 µs, and 16 heads of 128 goes from 222 µs to 241 µs. Qwen3-0.6B normalizes q and k in each of its 28 layers. If the measured deltas add up, that is roughly 1 ms per token at 16 threads and 0.1 ms at 4. Please report Qwen3-0.6B (8da4w) decode tokens/s before and after this PR, at 16 threads and at the runner's default. I'm asking for the data, not for the size gate back. If the loss is measurable, the opt-out above also serves as the mitigation.

Specific Comments

backends/xnnpack/operators/op_mean_dim.py:64,72 - the input is fetched twice (cast(torch.fx.Node, node.args[0]) and get_input_node(node, 0)), and the keepdim parse at :80 repeats the one in is_global_avg_pool_mean.
backends/xnnpack/operators/op_mean_dim.py:77 - this fires only when the layout pass breaks its own invariant (as in the first finding), but it reads like a user error. Say it is an internal layout error.
backends/xnnpack/test/ops/test_mean_dim.py:57-58 - "every other mean is one XNNStaticMean on an NCHW input, with no transpose" is wrong for the conv cases, which expect two XNNStaticTranspose. Limit the docstring to means on graph inputs.
docs/source/backends/xnnpack/op-support.csv:32 - "no dtype" can read as "any dtype"; line 2 says "no dtype conversion". Use "no dtype argument".
.ci/scripts/test_lora.sh:246,250 and .ci/scripts/test_lora_multimethod.sh:159 - Expected result prefix: echoes only the <think> form (:182, :115). The regex also accepts me\n and repeated <think> lines, so the log shows a prefix the check doesn't require. On failure, matches_base_response_prefix already prints the pattern.

Fix this →

Recommendation

Request Changes

The quantized channels_last lowering failure needs a fix and a test in this PR, or a test showing it doesn't happen. The opt-out and the end-to-end QK-norm number need discussion before merging but don't block on their own. The test and nit items don't block.

## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` as a convert to fp32, the fp32 `XNNStaticMean` from pytorch#24741, and a convert back to bf16, all inside the delegate. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - The converts are written into the program rather than added by the runtime at load time, as bf16 matmul results are. That keeps the change in Python and matches eager bf16 exactly, but a program exported now keeps the converts if XNNPACK's `xnn_define_static_reduce` later accepts bf16, which is the only thing that rejects a bf16 mean today. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` goes through the same converts as any other bf16 mean. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from pytorch#24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from pytorch#24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

JakeStevens · 2026-11-09T19:57:04Z

Thanks for the review. Replies to the review above, as of aa730da:

Cut-off: dropped. Small means are lowered like any other op the partitioner takes, and the description gives the thread-count cost, including end-to-end Qwen3 decode numbers (see the reply to the second review).
KeyError / wrong axes in ChannelsLastTaggedReshapePass: is_global_avg_pool_mean checks keepdim first and reads the rank from the mean's own meta["val"]. The visitor now refuses an NHWC input to a static mean, as an internal layout error, instead of reducing the wrong axes.
Tests assert the serialized node: Tester.check_xnnpack_node_count (with xnnpack_node_types moved from the layout pass tests into the Tester) checks for XNNGlobalAvgPooling2d vs XNNStaticMean, and for transposes, in the mean tests.
Dynamic tests run 5 random shapes.
Added: a qs8 convolution followed by a mean that is not Global Average Pooling, and a 6-D mean.
LoRA regex: pinned to the two observed continuations (assistant:\n<think>\n, repeated <think> allowed, and assistant: me\n). The comment names the cause (8da4w rounding once RMSNorm and QK-norm run in XNNPACK), and the scripts print the pattern on a mismatch.
Forward compatibility: the description says programs with the new node fail to load on older runtimes, including released prebuilt ones, with Unhandled node type, so the release note should say so.
NHWC copy for non-GAP means after a convolution: left as a follow-up, noted in the description.
Naming, why message, CSV wording: the cut-off and its why are gone. The CSV row now reads "no dtype argument".
xnn_define_static_reduce v1 vs v2: at the pinned XNNPACK, v1 is not deprecated. It copies the size_t axes to int64_t and calls v2, so the two behaved the same for our sorted non-negative axes. The runtime now calls xnn_define_static_reduce_v2, like the CPU provider.

Authored with AI using Claude Code.

JakeStevens · 2026-11-09T19:57:05Z

Replies to the second review, as of aa730da:

Quantized channels_last input. The exact graph (input → q → dq → mean(dq, [1], keepdim=True)) lowers and runs, because the input quantize stays outside the delegate, so the look-through is never taken. Once the quantize is inside the delegate (for example mean(x + x, [1])), the look-through reads meta["val"][0] as you traced. The mean was then delegated without a layout copy and reduced the wrong axes (here a runtime shape error; with other shapes it could be silent). The same misread made a qs8 convolution of a channels_last input after an add get a transpose that fails at runtime, on main too. is_nhwc_node and is_nchw_node now read the input's own XNN_NHWC_NODE tag. test_qs8_mean_dim_channels_last_input and test_qs8_channels_last_input_add_conv cover both, and both fail without the fix.

Opt-out for older runtimes. Not adding one: requiring a runtime with this change is acceptable here. The description names the Unhandled node type load error for the release note.

Tests.

test_conv_mean_meets_nhwc covers both LayerNorm2d paths, conv → y - y.mean(1, keepdim=True) and conv → mean → conv, in fp32 and qs8. In qs8 the mean can round one step (about 0.016) away from the reference, and that reaches the output, so the qs8 cases compare with atol=2e-2.
test_fp32_rms_norm checks XNNStaticMean: 1, no Global Average Pooling, and no edge pow, rsqrt, add or mul outside the delegate.
The qs8 mean tests check two XNNConvert (the input quantize and the output dequantize).

End-to-end QK-norm cost. Qwen3-0.6B 8da4w decode, after over before (above 1 is slower), executor_runner with the LLM custom ops. Qwen3 has 113 means here, and the delegates go from 368 to 283.

Device	1 thread	Runner default
Apple M-series (12P + 4E)	0.98x	1.11x to 1.12x (16 threads, 97 to 87 tokens/s)
iPhone 17	1.00x	1.05x (101 to 96 tokens/s)
Pixel 10	0.98x	1.00x

The cost is the small multi-head QK-norm means waking the thread pool, in line with your estimate. The mitigation we're taking is in XNNPACK: run small multi-output reductions as a single task, so pthreadpool runs them inline. That fixes it for every reduction of this shape without a cut-off tuned to one thread count. The numbers are in the description.

Specific comments.

The input is fetched once. The keepdim parse stays: it's one line and is_global_avg_pool_mean is also called from the layout pass.
The guard message says it is an internal layout error.
The _xnnpack_mean_nodes docstring is scoped to means of contiguous graph inputs.
The CSV says "no dtype argument".
The LoRA scripts no longer echo a single expected prefix for the base model checks.

Authored with AI using Claude Code.

JakeStevens · 2026-11-09T20:27:54Z

@claude review effort=max

claude · 2026-11-09T20:28:39Z

Claude finished @JakeStevens's task in 36m 30s —— View job

Approve: I found no blockers in aa730da. What's left is three test gaps, a small cleanup in the layout pass, and a release note that should cover more than LLMs. The test gaps are unpinned layout copies (my trace gives conv → y - y.mean(1, keepdim=True) two more transposes than main), a dtype test that never reaches its check, and a node-count helper that passes on typos.

Detailed review

PR Review: #24741

Summary

This PR lowers every mean.dim other than 4-D Global Average Pooling (GAP) to a new XNNStaticMean on an NCHW input. It also makes ChannelsLastTaggedReshapePass read a dequantized input's layout from that input's tag. I found no blockers. The problems are test gaps, a leftover in the pass, and the scope of the release note.

I couldn't run Python in this job, and the XNNPACK submodule isn't checked out. Every finding comes from reading the code, and I traced the transpose counts below by hand without running them.

Code Quality

is_nchw_node now duplicates is_nhwc_node, and the misread this PR removed from both is still in input_to_nhwc (backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py:147-156, :545-547).
- Both look-throughs now read the same XNN_NHWC_NODE tag. is_nchw_node is therefore not is_nhwc_node(node) in every case, and its dq branch can go.
- input_to_nhwc still tests input_node.meta["val"][0]. For a 4-D placeholder that slice is 3-D, so _is_nhwc always returns False and the early return never fires.
- The real check is input_dim_order at :551, so these lines are dead rather than wrong. Read meta["val"] as input_to_nchw does at :624, or drop the branch.
  Fix this →

Testing

The tests don't pin the layout copies where this PR moves them. The description says the GAP .ptes are byte-identical to main and that the tests cover the layout copies the pass inserts. Three places don't check them:
- GAP. _xnnpack_mean_nodes (backends/xnnpack/test/ops/test_mean_dim.py:81) and the GAP case of test_fp32_conv_mean_dim (:257) don't count XNNStaticTranspose. Byte-identity depends on that count, and this PR changed both the mean routing and the dequantize look-through that decide it. My trace gives 2 for each fp32 case.
- A mean that meets NHWC data. test_conv_mean_meets_nhwc (:321) counts only XNNStaticMean.
  - For fp32 ConvSubtractMean, my trace gives 4 transposes in one delegate. Main has 2: one for the conv input and one at the delegate boundary. Main also runs the mean on the portable kernel and has a second delegate.
  - converge_layouts (channels_last_tagged_reshape_pass.py:251-261) makes the subtraction NHWC, because both operands are 4-D. That copies the mean's result to NHWC, then copies the full-size difference back to NCHW at the output.
  - So the description's "the same NHWC to NCHW copy main made at the delegate boundary" doesn't hold for this pattern.
  - ConvMeanConv gives 4, the same as main.
  - Pinning both counts also shows what the axis-remap follow-up saves.
- The pass test. test_fp32_channels_last_tagged_reshape_pass_conv_mean_channels (backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py:925) counts 2 _to_copy.
  - The old pass also inserts 2 here: the input to NHWC, and the NHWC mean output back to NCHW. So the test passes either way.
  - Assert instead that the mean's input is the contiguous _to_copy that the comment at :919 describes.
    Fix this →
test_fp32_mean_dim_unsupported_dtype never reaches the dtype check (test_mean_dim.py:327-335, generic_node_configs.py:725-727).
- Its fp64 output already fails _check_outputs_are_valid_dtypes (xnnpack_config.py:278) inside check_common_constraints. The test therefore still passes if the dtype check is removed.
- With this PR, that check is what keeps a mean like x.mean(-1, dtype=torch.float32) on an fp16 tensor on the portable kernel. Without it, the visitor's own check (op_mean_dim.py:55-58) raises inside to_edge_transform_and_lower.
- Use an fp16 input with dtype=torch.float32 on a non-GAP dim such as (-1,). Both dtypes pass the generic check, because _check_inputs_are_valid_dtypes only compares the inputs with each other.
  Fix this →
check_xnnpack_node_count passes on a misspelled node type (backends/xnnpack/test/tester/tester.py:175-181).
- xnnpack_node_types returns a Counter, so an unknown name reads as 0.
- Most of the new pins are zeros ("XNNStaticTranspose": 0, "XNNStaticMean": 0, "XNNGlobalAvgPooling2d": 0). A typo or a later rename would silently turn them into no-ops.
- Assert that each key is a member name of XNodeUnion (xnnpack_graph_schema.py:402).
  Fix this →

Backward Compatibility

The release note should cover every non-GAP mean, not only LLMs (generic_node_configs.py:715-736).
- The description ties the load failure on older runtimes to re-exported LLM programs (RMSNorm, QK-norm).
- But any XNNPACK export with a mean.dim other than a 4-D [2, 3] keepdim=True mean now contains XNNStaticMean. Examples are the x.mean([2, 3]) heads (keepdim=False) of torchvision's ShuffleNetV2 and MNASNet, which used to stay on the portable kernel.
- State that condition in the note, together with the Unhandled node type error, so users of non-LLM models see it too.

Recommendation

Approve

Nothing blocks. The transpose pins, the dtype test and the key check are small test changes, and the release note is a change to the description. The is_nchw_node cleanup is optional.

## Ulterior Motive Run RMSNorm in LeRobot policies on XNNPACK (pytorch#24667), where `mean.dim` fails to partition with "only supports averaging 4D tensors, got tensor of rank 3". ## Rationale **What**: Add an `XNNStaticMean` node backed by `xnn_define_static_reduce(xnn_reduce_mean)`, and lower every `mean.dim` that is not Global Average Pooling to it: ranks 1 to 6, any dims, and either `keepdim`. **Why**: XNNPACK only lowered a 4D mean over dims [2, 3] with `keepdim=True`. RMSNorm averages the last dim of a 3-D tensor, so its mean stayed on the portable kernel and split the delegate around every norm. ## Details Review order: the serialization schema and `XNNCompiler.cpp`, then `op_mean_dim.py` and `is_global_avg_pool_mean`, then `ChannelsLastTaggedReshapePass`, then `MeanDimConfig`, then the tests. - The 4D, dims [2, 3], `keepdim=True` case still lowers to Global Average Pooling on NHWC. For those programs the `.pte` is byte-identical to main (fp32, fp16, qs8, conv-bn-hardtanh-mean, and a CNN ending in `AdaptiveAvgPool2d(1)`), and the tests check which XNNPACK node each mean serializes to and how many transposes surround it. - Every other `mean.dim` reduces over NCHW dims, so `ChannelsLastTaggedReshapePass` gives it an NCHW input, and the visitor refuses an NHWC one as an internal layout error. After a convolution that adds layout copies inside the one delegate. For conv then `mean(1, keepdim=True)` then conv, the program has 4 transposes, as on main. For `y - y.mean(1, keepdim=True)` after a convolution (the centering in ConvNeXt's LayerNorm2d), the subtraction stays NHWC, so its full-size result is copied back to NCHW at the output: 4 transposes in one delegate, where main has 2 in two delegates around a portable mean. At 96 channels and 56x56 that is 0.75x at 1 thread, 0.85x at 4 and 1.12x at 16, and conv, mean, conv is 0.54x, 0.70x and 0.90x. Remapping the reduction axes to read NHWC directly would remove these copies, also for heads like `x.mean([2, 3])`; that is left for a follow-up. - `is_global_avg_pool_mean` takes the rank from the mean's own output, which with `keepdim=True` equals the input's, because inside `ChannelsLastTaggedReshapePass` the mean's input can be a newly inserted node without `"val"`. - The node uses a shared `_XNNStaticReduce` table, so other XNNPACK reductions can reuse it and `defineGenericReduceNode`. It is appended to `XNodeUnion`, so existing node IDs do not change. Practically every re-exported XNNPACK LLM program contains it, because RMSNorm (`examples/models/llama/norm.py`) and QK-norm lower to it, and those programs fail to load on runtimes built before this change, including released prebuilt runtimes, with `Unhandled node type` in the log. The release note should say so. - `ChannelsLastTaggedReshapePass` now reads the layout of a quantized graph input from the input's own tag. It used to test `meta["val"][0]`, a 3-D slice that is never channels_last, so a channels_last input behind a quantize and dequantize read as NCHW. A mean over NCHW dims of such an input then reduced the wrong axes, and on main a convolution of it got a transpose that fails at runtime. Both now lower correctly. - An empty dim list averages over every dim, as in PyTorch. - Latency against main, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: RMSNorm 1x128x2048 0.48x, 3-D last-dim mean 0.45x, RMSNorm then Linear 0.77x, 4 pre-norm residual MLP blocks 0.99x for decode (1 token) and 0.80x for prefill (64 tokens), conv then mean over channels 0.88x. - XNNPACK splits a reduction with more than one output into threadpool tasks, so a small mean with several outputs can cost more in XNNPACK than on the portable kernel when the pool has many threads. For RMSNorm on 10 shapes from 1x1x128 to 1x32x2048, every shape is 0.43x to 0.97x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.55x (1.5 to 2.7 us more), and the rest 0.44x to 0.93x. At 16 threads those small shapes are 3.7x to 5.7x (21 to 23 us more), and the rest 0.70x to 0.98x. Inside a model with neighbouring delegated ops the cost can remain: decode QK-norm between two linears, 8 heads of 64, goes from 11.8 us to 39.1 us at 16 threads, and 16 heads of 128 from 222 us to 241 us. A mean alone in its delegate pays the same wake-up (a bare 1x128x128 mean is 3.26x at 16 threads and 1.14x at 4), as does any op lowered alone, such as `add` (2.95x at 1x128x128 and 16 threads on main). These means are lowered like every other op the XNNPACK partitioner takes, rather than gated by size, which would be tuned to one thread count and treat static and dynamic exports differently. End to end, Qwen3-0.6B with 8da4w quantization (113 RMSNorm and QK-norm means, 368 delegates on main and 283 here) decodes, after over before in `executor_runner` with the LLM custom ops: on the same Apple M-series CPU 0.98x at 1 thread and 1.11x to 1.12x at the default 16 threads (97 to 87 tokens/s); on an iPhone 17 1.00x at 1 thread and 1.05x at the runner's default (101 to 96 tokens/s); on a Pixel 10 0.98x and 1.00x. The cost is the multi-head QK-norm means. Running small multi-output reductions without waking the pool belongs in XNNPACK and is left for a follow-up. - The LoRA CI scripts check the start of the 8da4w-quantized Qwen3 base model's greedy output with a regex. With RMSNorm and QK-norm now in XNNPACK, its first token can tip from `:\n` to `:` (`assistant: me` instead of `assistant:\n<think>`): unquantized, main and this branch generate the same tokens with logits within 1e-4, but 8da4w re-rounds activations and the two candidates are within about 0.5 logits. The regex accepts exactly those two continuations, and the scripts print it when it does not match. - Tests cover 1-D to 6-D inputs, all dims, `keepdim=False`, fp16 and qs8 over the innermost and a middle dim, dynamic kept and reduced dims over 5 random shapes, RMSNorm, fp32 and qs8 convolutions followed by a mean, a mean over channels whose result meets the convolution output or feeds another convolution (as in ConvNeXt's LayerNorm2d), qs8 channels_last inputs to a mean and to a convolution, the serialized XNNPACK nodes (moved `xnnpack_node_types` from the layout pass tests to the XNNPACK `Tester`), and the layout copies the pass inserts. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from pytorch#24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from pytorch#24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details - `defineGenericReduceNode` adds the converts when a serialized value is bf16, the same way `defineSoftmaxNode` does since pytorch#24805. It reuses that change's helpers: `convertIfDistinct`, and `defineBf16SoftmaxFp32Value`, renamed `defineFp32ValueForBf16` now that reductions use it too. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to main. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in #24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from #24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details - `defineGenericReduceNode` adds the converts when a serialized value is bf16, the same way `defineSoftmaxNode` does since #24805. It reuses that change's helpers: `convertIfDistinct`, and `defineBf16SoftmaxFp32Value`, renamed `defineFp32ValueForBf16` now that reductions use it too. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in #24741. fp32 and fp16 programs are byte-identical to main. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in #24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in #24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which #24745 lowers for programs exported with dim order, and #24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.

JakeStevens requested review from digantdesai and mergennachin as code owners November 6, 2026 18:03

JakeStevens added the release notes: xnnpack Changes to the XNNPack backend delegate label Nov 6, 2026

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 6, 2026

JakeStevens requested a review from shoumikhin November 6, 2026 19:20

JakeStevens force-pushed the xnnpack-mean-dim-static-reduce branch from 7e905e0 to 99c4063 Compare November 6, 2026 20:06

JakeStevens mentioned this pull request Nov 6, 2026

Lower bf16 mean.dim to XNNPACK as an fp32 static mean #24746

Merged

JakeStevens mentioned this pull request Nov 6, 2026

[XNNPACK] Lower fp32 <-> bf16 casts to XNNPACK converts #24745

Merged

JakeStevens force-pushed the xnnpack-mean-dim-static-reduce branch from caf6ec8 to ac2756c Compare November 9, 2026 17:10

JakeStevens force-pushed the xnnpack-mean-dim-static-reduce branch from ac2756c to 9b5b8da Compare November 9, 2026 17:49

JakeStevens force-pushed the xnnpack-mean-dim-static-reduce branch from 9b5b8da to aa730da Compare November 9, 2026 19:40

JakeStevens force-pushed the xnnpack-mean-dim-static-reduce branch from aa730da to e578110 Compare November 9, 2026 21:52

shoumikhin approved these changes Nov 9, 2026

View reviewed changes

JakeStevens merged commit b3e2457 into pytorch:main Nov 10, 2026
389 of 390 checks passed

JakeStevens deleted the xnnpack-mean-dim-static-reduce branch November 10, 2026 02:08

Conversation

JakeStevens commented Nov 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Ulterior Motive

Rationale

Details

Uh oh!

pytorch-bot Bot commented Nov 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24741

❌ 1 New Failure

Uh oh!

JakeStevens commented Nov 6, 2026

Uh oh!

claude Bot commented Nov 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24741

Summary

Code Quality

Testing

Backward Compatibility

Performance

Specific Comments

Recommendation

Uh oh!

JakeStevens commented Nov 9, 2026

Uh oh!

claude Bot commented Nov 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24741

Summary

Backward Compatibility

Testing

Performance

Specific Comments

Recommendation

Uh oh!

JakeStevens commented Nov 9, 2026

Uh oh!

JakeStevens commented Nov 9, 2026

Uh oh!

JakeStevens commented Nov 9, 2026

Uh oh!

claude Bot commented Nov 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24741

Summary

Code Quality

Testing

Backward Compatibility

Recommendation

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

JakeStevens commented Nov 6, 2026 •

edited

Loading

pytorch-bot Bot commented Nov 6, 2026 •

edited

Loading

claude Bot commented Nov 6, 2026 •

edited

Loading

claude Bot commented Nov 9, 2026 •

edited

Loading

claude Bot commented Nov 9, 2026 •

edited

Loading